import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from linear_regression import LinearRegression

data = pd.read_csv('../data/non-linear-regression-x-y.csv')

# Fetch traingin set and labels.
x = data['x'].values.reshape((data.shape[0], 1))
y = data['y'].values.reshape((data.shape[0], 1))

# Print the data table.
data.head(10)

plt.plot(x, y)
plt.show()

input_param_name = 'Economy..GDP.per.Capita.'
output_param_name = 'Happiness.Score'

# Set up linear regression parameters.
num_iterations = 50000  # Number of gradient descent iterations.
regularization_param = 0  # Helps to fight model overfitting.
learning_rate = 0.02  # The size of the gradient descent step.
polynomial_degree = 15  # The degree of additional polynomial features.
sinusoid_degree = 15  # The degree of sinusoid parameter multipliers of additional features.
normalize_data = True  # Flag that indicates that data needs to be normalized before training.

# Init linear regression instance.
linear_regression = LinearRegression(x, y, polynomial_degree, sinusoid_degree, normalize_data)

# Train linear regression.
(theta, cost_history) = linear_regression.train(
    learning_rate,
    # regularization_param = 0  # Helps to fight model overfitting.
    num_iterations
)

# Print training results.
print('Initial cost: {:.2f}'.format(cost_history[0]))
print('Optimized cost: {:.2f}'.format(cost_history[-1]))

# Print model parameters
theta_table = pd.DataFrame({'Model Parameters': theta.flatten()})
theta_table

# Plot gradient descent progress.
plt.plot(range(num_iterations), cost_history)
plt.xlabel('Iterations')
plt.ylabel('Cost')
plt.title('Gradient Descent Progress')
plt.show()

# Get model predictions for the trainint set.
predictions_num = 1000
x_predictions = np.linspace(x.min(), x.max(), predictions_num).reshape(predictions_num, 1);
y_predictions = linear_regression.predict(x_predictions)

# Plot training data with predictions.
plt.scatter(x, y, label='Training Dataset')
plt.plot(x_predictions, y_predictions, 'r', label='Prediction')
plt.show()